from PySide6 import QtCore, QtGui, QtWidgets

import os
import sys
import subprocess
import time

from ui_infodialog import Ui_InfoDialog
from ui_mainwindow import Ui_MainWindow
import fire

from typing import List
import requests
import json

server : str = 'http://127.0.0.1:11434'
# server : str = 'http://192.168.1.130:11434'
# server : str = 'http://118.178.197.102:11434'
qwen_model : str = 'qwen2.5:0.5b'

from ollama.api import *
class OllamaApi:

    def __init__(self):
        pass

    def isValid(self) -> bool:
        '''检查ollama服务是否启动'''
        try:
            response = requests.get(f'{server}/api/version')
            if response.status_code == 200:
                print(response.content)
                return True
        except (requests.exceptions.ConnectionError) as e:
            print(e)
            # 可能服务器没有响应，请尝试其他IP或端口
        return False
    
    def version(self, timeout=0.1) -> Version:
        resp = requests.get(f'{server}/api/version', timeout=timeout)
        print(resp)
        return Version(orginal=resp.json(), **resp.json())

    def tags(self) -> Tags:
        '''列出本地模型'''
        resp = requests.get(f'{server}/api/tags')
        print(resp)
        return Tags(orginal=resp.json(), **resp.json())
        
    def show(self, model_name) -> Show:
        '''显示模型信息'''
        resp = requests.post(f'{server}/api/show', json={"model": model_name})
        print(resp)
        return Show(orginal=resp.json(), **resp.json())

    def ps(self) -> Ps:
        '''列出正在运行的模型'''
        return Ps(**requests.get(f'{server}/api/ps').json())


class Chatbot():
    
    # def __init__(self, sever='http://127.0.0.1:11434', qwen_model='qwen2.5:0.5b'):
    #     self.server = sever
    #     self.qwen_model = qwen_model
    #     self.response = None

    def __init__(self):
        self.response = None

    def curl(self, prompt:str):
        url = server + '/api/generate'
        data = {
            'model': qwen_model,
            'prompt': prompt
        }
        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.post(url, data=json.dumps(data), headers=headers, stream=True, allow_redirects=True)
        return response

    def generate(self, prompt:str):
        try:
            self.response = self.curl(prompt)
        except Exception as e:
            # yield str(e)
            yield '请求异常，请检查网络连接.'
            return

        # 检查响应状态码
        if self.response.status_code!= 200:
            yield f'请求失败，状态码: {self.response.status_code}'
            # st.session_state.messages.append({'role': 'ai', 'content': errstr, 'error': True})
            # st.warning(errstr)
        else:
            reply = ''
            try:
                for line in self.response.iter_lines():
                    answer = json.loads(line.decode(errors='ignore'))
                    reply = answer['response']
                    # if reply != "":
                        # print(answer['response'])
                    if answer['done']:
                        break
                    yield reply
            finally:
                self.response.close()
                self.response = None

    def terminate(self):
        if self.response is not None:
            try:
                self.response.close()
                self.response.connection
            except:
                self.response = None

chatbot = Chatbot()

class OllamaApiManager(QtCore.QObject):

    version = QtCore.Signal()
    versionSignal = QtCore.Signal(Version, bool)

    tags = QtCore.Signal()
    tagsSignal = QtCore.Signal(Tags, bool)

    show = QtCore.Signal(str)
    showSignal = QtCore.Signal(Show, bool)

    def __init__(self, parent: QtCore.QObject=None):
        super().__init__(parent)

        self.m_thread = QtCore.QThread()

    def stop(self):
        if self.m_thread.isRunning():
            self.m_thread.quit()
            self.m_thread.wait()

    def start(self):
        self.stop()

        self.worker = OllamaApiWorker(self)
        self.version.connect(self.worker.version)
        self.tags.connect(self.worker.tags)
        self.show.connect(self.worker.show)

        self.m_thread.destroyed.connect(self.worker.deleteLater)
        self.worker.moveToThread(self.m_thread)
        self.m_thread.start()

class OllamaApiWorker(QtCore.QObject):

    def __init__(self, manager: OllamaApiManager):
        super(OllamaApiWorker, self).__init__()

        self.manager = manager
        self.ollama_api = OllamaApi()

    def version(self):
        try:
            v = self.ollama_api.version()
            self.manager.versionSignal.emit(v, True)
        except:
            v = ""
            self.manager.versionSignal.emit(None, False)

    def tags(self):
        try:
            tags = self.ollama_api.tags()
            self.manager.tagsSignal.emit(tags, True)
        except:
            self.manager.tagsSignal.emit(None, False)

    def show(self, model):
        try:
            show = self.ollama_api.show(model)
            self.manager.showSignal.emit(show, True)
        except:
            self.manager.showSignal.emit(None, False)

class Worker(QtCore.QThread):
    Output = QtCore.Signal(str)
    
    def __init__(self, parent = ...):
        super().__init__(parent)
        self.elapsed = 0
        self.token_in_sec = 0

    def submit(self, text):
        self.text = text
        self.start_time = time.monotonic()
        self.token_in_sec = 0

    def run(self):
        chatbot.terminate()
        channel = chatbot.generate(self.text)
        response = ""
        tokens = 0
        for chunk in channel:
            if isinstance(chunk, str):
                response += chunk
                tokens += 1
                self.elapsed = time.monotonic() - self.start_time
                self.token_in_sec = tokens / self.elapsed
                self.Output.emit(response)

def random_say():
    import random
    says = [
        '你是谁',
        '1+1等于几',
        '你能讲个笑话吗',
        '给我来个有趣的故事吧',
        '你能描述一下你最喜欢的电影吗',
        '你能告诉我一个有趣的事实吗',
        '你能写一段诗吗',
        '你能创作一段音乐吗',
        '你能画一幅画吗',
        '你能推荐一部好书吗',
        '你能给我来个有趣的谜语吗',
    ]
    return random.choice(says)


class InfoDialog(QtWidgets.QDialog):
    model :str = ""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.ui = Ui_InfoDialog()
        self.ui.setupUi(self)

    def setModel(self, model: str):
        self.model = model
        self.setWindowTitle(f'模型信息：{model}')

    def setModelShow(self, show:Show):
        from dataclasses import asdict
        self.ui.model_show_editor.setPlainText(json.dumps(show.getOrginal(), indent=4)) 
        self.setModelfile(show.modelfile)

        if 'system' in show.getOrginal():
            self.ui.model_system_editor.setPlainText(show.system)
        else:
            self.ui.tab_system.setHidden(True)


    def setModelfile(self, modelfile):
        self.ui.model_modelfile_editor.setPlainText(modelfile)
        

class OllamaWindow(QtWidgets.QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.api = OllamaApi()
        self.ui.edit_server.setText(server)
        self.ui.comb_models.setSizeAdjustPolicy(QtWidgets.QComboBox.SizeAdjustPolicy.AdjustToContents)

        self.ui.splitter.setSizes([700,400])
        self.ui.splitter_2.setSizes([500, 100])

        self.worker = Worker(self)
        self.worker.Output.connect(self.update_output)

        self.ui.btn_send.clicked.connect(self.do_question)
        self.ui.btn_random.clicked.connect(self.do_question)
        self.ui.btn_simple.clicked.connect(self.do_question)
        self.ui.btn_deep_analysis.clicked.connect(self.do_question)
        self.ui.btn_translator.clicked.connect(self.do_question)
        self.ui.btn_what_this.clicked.connect(self.do_question)
        self.ui.btn_copy.clicked.connect(lambda : self.clipboard.setText(self.ui.text_output.toMarkdown()))
        self.ui.btn_copy_pixmap.clicked.connect(self.on_btn_copy_pixmap_clicked)
        
        self.ui.btn_c.clicked.connect(self.do_question)
        self.ui.btn_cpp.clicked.connect(self.do_question)
        self.ui.btn_python.clicked.connect(self.do_question)
        self.ui.btn_golang.clicked.connect(self.do_question)
        
        # manager
        self.ollama_api_manager = OllamaApiManager(self)
        self.ollama_api_manager.versionSignal.connect(self.on_ollama_api_versionSignal)
        self.ollama_api_manager.tagsSignal.connect(self.on_ollama_api_tagsSignal)
        self.ollama_api_manager.showSignal.connect(self.on_ollama_api_showSignal)
        self.ollama_api_manager.start()
        
        # init/options
        self.version()
        self.init_clipboard_options()

        self.ui.text_input.setAcceptRichText(False)
        self.ui.text_output.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.ui.text_output.customContextMenuRequested.connect(self.on_text_output_customContextMenuRequested)

        # self.ui.btn_model_info.connect(self.on_btn_model_info_clicked) # auto connected

        
    def pixmap_merge_v3(self, pixmap1, pixmap2, pixmap3, vertical=False):
        return self.pixmap_merge(self.pixmap_merge(pixmap1, pixmap2, vertical), pixmap3, vertical, True)

    def pixmap_merge(self, pixmap1: QtGui.QPixmap, pixmap2: QtGui.QPixmap, vertical=False, right=False):
        # horizontal
        width_h = pixmap1.width() + pixmap2.width() + 6
        height_h = pixmap1.height() > pixmap2.height() and pixmap1.height() or pixmap2.height()

        # vertical
        width_v = pixmap1.width() > pixmap2.width() and pixmap1.width() or pixmap2.width()
        height_v = pixmap1.height() + pixmap2.height() + 6

        # note: h in vertical, v in horizontal
            # h:x0y1,x1y1
            # v:x1y0,x1y1        
        line_rect = QtCore.QRect()
        if vertical:
            line_rect.setTopLeft(QtCore.QPoint(0, pixmap1.height() +2))
            line_rect.setTopRight(QtCore.QPoint(width_v, pixmap1.height() +2))
        else:
            line_rect.setTopLeft(QtCore.QPoint(pixmap1.width() + 2, 0))
            line_rect.setBottomLeft(QtCore.QPoint(pixmap1.width() + 2, height_h))

        if vertical:
            pixmap = QtGui.QPixmap(width_v, height_v)
        else:
            pixmap = QtGui.QPixmap(width_h, height_h)
        
        pixmap.fill(QtCore.Qt.GlobalColor.lightGray)
        painter = QtGui.QPainter(pixmap)
        
        if vertical:
            painter.drawLine(line_rect.topLeft(), line_rect.topRight())
            painter.drawPixmap(QtCore.QPoint(0,0), pixmap1)
            if right:
                painter.drawPixmap(QtCore.QPoint(pixmap1.width() - pixmap2.width(),pixmap1.height()+6), pixmap2)
            else:    
                painter.drawPixmap(QtCore.QPoint(0,pixmap1.height()+6), pixmap2)
        else:
            painter.drawLine(line_rect.topLeft(), line_rect.bottomLeft())
            painter.drawPixmap(QtCore.QPoint(0,0), pixmap1)
            painter.drawPixmap(QtCore.QPoint(pixmap1.width() + 6,0), pixmap2)
        painter.end()

        return pixmap

    def server_pixmap(self):
        text = '服务器地址: ' + self.ui.edit_server.text()
        server_metrics = QtGui.QFontMetrics(self.font())
        server_pixmap = QtGui.QPixmap(server_metrics.size(QtGui.Qt.TextFlag.TextSingleLine , text))
        server_pixmap.fill(QtCore.Qt.GlobalColor.transparent)
        server_painter = QtGui.QPainter(server_pixmap)
        server_painter.drawText(self.ui.edit_server.rect(), text)
        server_painter.end()
        return server_pixmap

    def on_btn_copy_pixmap_clicked(self):
        '''复制为图片'''

        iw = self.ui.text_input.document().size().width()
        ih = self.ui.text_input.document().size().height()
        ow = self.ui.text_output.document().size().width()
        oh = self.ui.text_output.document().size().height()
       
        # final size
        input_size = QtCore.QSize(iw, ih)
        output_size = QtCore.QSize(ow, oh)

        # cache input pixmap
        input_pixmap = QtGui.QPixmap(input_size)
        input_pixmap.fill(QtCore.Qt.GlobalColor.transparent)
        input_painter = QtGui.QPainter(input_pixmap)
        self.ui.text_input.document().drawContents(input_painter)
        input_painter.end()
        
        # cache output pixmap
        output_pixmap = QtGui.QPixmap(output_size)
        output_pixmap.fill(QtCore.Qt.GlobalColor.transparent)
        output_painter = QtGui.QPainter(output_pixmap)
        self.ui.text_output.document().drawContents(output_painter)
        output_painter.end()

        # cache model pixmap
        model_metrics = QtGui.QFontMetrics(self.ui.comb_models.font())
        model_pixmap = QtGui.QPixmap(model_metrics.size(QtGui.Qt.TextFlag.TextSingleLine , '模型:' + self.ui.comb_models.currentText()))
        model_pixmap.fill(QtCore.Qt.GlobalColor.transparent)
        model_painter = QtGui.QPainter(model_pixmap)
        model_painter.drawText(self.ui.comb_models.rect(), '模型:' + self.ui.comb_models.currentText())
        model_painter.end()

        # cache token pixmap
        token_metrics = self.ui.statusbar.fontMetrics()
        token_pixmap = QtGui.QPixmap(token_metrics.size(QtGui.Qt.TextFlag.TextSingleLine, self.ui.statusbar.currentMessage()))
        token_pixmap.fill(QtCore.Qt.GlobalColor.transparent)
        token_painter = QtGui.QPainter(token_pixmap)
        token_painter.drawText(self.ui.statusbar.rect(), self.ui.statusbar.currentMessage())
        token_painter.end()

        menu = QtWidgets.QMenu(self)
        menu.addAction('左右结构', lambda: self.clipboard.setPixmap(self.pixmap_merge(input_pixmap, self.pixmap_merge(output_pixmap, self.pixmap_merge(model_pixmap, token_pixmap), True, True), False)))
        menu.addAction('上下结构', lambda: self.clipboard.setPixmap(self.pixmap_merge_v3(input_pixmap, output_pixmap, self.pixmap_merge(model_pixmap, token_pixmap), True)))
        menu.addAction('上下结构+服务器地址', lambda: self.clipboard.setPixmap(self.pixmap_merge(self.pixmap_merge_v3(input_pixmap, output_pixmap, self.server_pixmap(), True), self.pixmap_merge(model_pixmap, token_pixmap), True, True)))
        menu.exec(QtGui.QCursor.pos())

    def on_text_output_customContextMenuRequested(self, pos):
        menu = QtWidgets.QMenu(self)
        menu.addAction('复制', lambda: self.clipboard.setText(self.ui.text_output.toMarkdown()))
        menu.addAction('复制为图片', self.on_btn_copy_pixmap_clicked)
        menu.exec(self.ui.text_output.mapToGlobal(pos))

    def on_ollama_api_versionSignal(self, v: Version, b):
        '''响应Ollama API版本'''
        if b:
            self.setWindowTitle('Ollama 测试器 - 服务器版本 v' + v.version)
            self.ollama_api_manager.tags.emit() # 有版本了就请求模型列表
        else:
            self.ui.statusbar.showMessage('Ollama API: 请求模型版本出错')

    def on_ollama_api_tagsSignal(self, tags: Tags, b):
        '''响应Ollama API标签'''
        self.ui.comb_models.clear()
        if b:
            self.ui.comb_models.addItems([model.name for model in tags.models])
            
            model = min(tags.models, key=lambda x: x.size) # 默认选中最小大小的模型
            self.ui.comb_models.setCurrentText(model.name)
            self.welcome() # 有模型了就请求 welcome
        else:
            self.ui.statusbar.showMessage('Ollama API: 请求模型列表出错')

    def on_ollama_api_showSignal(self, show: Show, b):
        '''响应Ollama API信息'''
        if b:
            dialog = InfoDialog()
            dialog.setModel(self.ui.comb_models.currentText())
            dialog.setModelShow(show)
            dialog.exec()
        else:
            self.ui.statusbar.showMessage('Ollama API: 请求模型信息出错')

    def welcome(self):
        '''欢迎'''
        self.worker.submit("你是一只可爱的猫娘，请用50字以内的欢迎信息迎接我（例如：你好！‘我是一只可爱的小猫咪，很高兴为您服务。有什么我可以帮忙的吗？‘），语气要可爱、温暖，保持猫娘的身份，不要反转角色。")
        QtCore.QTimer.singleShot(1, self.worker.start)

    def version(self):
        '''版本'''
        self.ollama_api_manager.version.emit()

    def init_clipboard_options(self):
        '''初始化剪贴板选项'''
        self.ui.comb_clipboard_action.addItem('总结内容')
        self.ui.comb_clipboard_action.addItem('翻译')
        
        self.clipboard = QtWidgets.QApplication.clipboard()
        self.clipboard.dataChanged.connect(self.on_clipboard_changed)

    def on_clipboard_changed(self):
        '''剪贴板有变化时'''
        if self.ui.comb_clipboard_action.currentIndex() == 0:
            return
        
        if self.clipboard.mimeData().hasText():
            self.ui.text_input.setPlainText(self.clipboard.text())
            # self.ui.btn_translator.clicked.emit()
            question = self.ui.text_input.toPlainText()
            question += f'\n\n 「{self.ui.comb_clipboard_action.currentText()}」'
            self.worker.submit(question)
            
            QtCore.QTimer.singleShot(1, self.worker.start)

    @QtCore.Slot()
    def on_action_config_server_triggered(self):
        print("config server")

    @QtCore.Slot(str)
    def on_edit_server_textChanged(self, text: str):
        global server
        server = text
        self.version()
        
    @QtCore.Slot(str)
    def on_comb_models_currentTextChanged(self, text: str):
        if text == '':
            return
        
        global qwen_model
        qwen_model = text

    def update_output(self, output: str):
        self.ui.raw_output.setPlainText(output) # 保存原始输出到 raw output
        self.ui.raw_output.verticalScrollBar().setValue(self.ui.raw_output.verticalScrollBar().maximum())

        if output.startswith("<think>"):        # 对于 <think> </think> 的情况，采用
            output = output.replace("<think>", "Thinking:")
            if output.find("</think>").__ne__(-1):
                output = output.replace("</think>", "\n\n----\n")
            else:
                output = output + "\n\n----\n"
        self.ui.text_output.setMarkdown(output)
        self.ui.text_output.verticalScrollBar().setValue(self.ui.text_output.verticalScrollBar().maximum())
        self.ui.statusbar.showMessage(f'耗时：{self.worker.elapsed:.2f}s ({self.worker.token_in_sec:.1f} token/s)')

    def do_question(self):
        btn: QtWidgets.QPushButton = self.sender()
        self.worker.terminate()

        if btn == self.ui.btn_send:
            self.worker.submit(self.ui.text_input.toPlainText())
            QtCore.QTimer.singleShot(1, self.worker.start)
            return

        if btn == self.ui.btn_random:
            question = random_say()
            self.ui.text_input.setPlainText(question)
            self.worker.submit(question)
            QtCore.QTimer.singleShot(1, self.worker.start)
            return

        question = self.ui.text_input.toPlainText()
        question += '\n\n' + btn.text()
        self.worker.submit(question)
        
        QtCore.QTimer.singleShot(1, self.worker.start)

    @QtCore.Slot(int)
    def on_comb_clipboard_action_currentIndexChanged(self, index):
        if index == 0:
            self.ui.groupBox_3.setVisible(True)
            return
        
        self.ui.groupBox_3.setVisible(False)

    @QtCore.Slot()
    def on_btn_model_info_clicked(self):
        self.ollama_api_manager.show.emit(self.ui.comb_models.currentText())
        
if __name__ == '__main__':

    if len(sys.argv) > 1:
        fire.Fire(OllamaApi)
        pass
    else:
        app = QtWidgets.QApplication(sys.argv)
        widget = OllamaWindow()
        widget.setWindowTitle("Ollama 测试器")
        widget.show()
        widget.move(QtGui.QGuiApplication.primaryScreen().availableGeometry().center() - widget.rect().center())
        sys.exit(app.exec())
